from turtle import forward
from nets.bases.basicmodel import BasicModule
import torch as t
import torch.nn as nn
import torch.nn.functional as F
from nets.modules.depthwise_block import DepthWiseConv

class DownSample(nn.Module):
    
    def __init__(self, inplace, outplace, isDown = True):
        super(DownSample, self).__init__()
        
        self.layer1 = nn.Sequential(
            DepthWiseConv(inplace, inplace),
            nn.BatchNorm2d(inplace),
            nn.ReLU(inplace = True),
            
            nn.Conv2d(inplace, outplace, (3, 3), stride=1, padding=1, bias=False),
            nn.BatchNorm2d(outplace),
            nn.ReLU(inplace = True),
        )
        
        self.isDown = True
        
        
    def forward(self, x):
        x = self.layer1(x)
        if self.isDown:
            x = F.max_pool2d(x, 2)
        return x
    

    

class UpSample(nn.Module):
    
    def __init__(self, in_plances, out_plances, isUp = True):
        super().__init__()

        self.model = nn.Sequential(
            nn.Conv2d(in_plances, out_plances, (3, 3), stride=1, padding=1, bias=False),
            nn.BatchNorm2d(out_plances),
            nn.ReLU(out_plances),
        )
        if isUp:        
            self.upSample = nn.Upsample(scale_factor=2, mode='bilinear',align_corners=False)
        else:
            self.upSample = nn.Conv2d(out_plances, out_plances, (1, 1), stride=1, padding=0, bias=False)
        
    
    def forward(self, x):
        x = self.model(x)
        x = self.upSample(x)
        return x
    

class CenterLayer(nn.Module):
    
    def __init__(self, in_plances, out_plances):
        super().__init__()
        
        self.model = nn.Sequential(
            nn.Conv2d(in_plances, out_plances, (3, 3), stride=1, padding=1, bias=False),
            nn.BatchNorm2d(out_plances),
            nn.ReLU(out_plances),
        )
        
    def forward(self, x):
        x = self.model(x)
        return x
    

class UNetModel(BasicModule):
    
    def __init__(self):
        super(UNetModel, self).__init__()
        
        self.model_name = 'UNetModel_Dilated'
        self.down_layer1 = DownSample(inplace = 3, outplace = 12, isDown = True) # 256
        self.down_layer2 = DownSample(inplace = 12, outplace = 126, isDown = True) # 128
        self.down_layer3 = DownSample(inplace = 126, outplace = 252, isDown = True) # 64
        self.down_layer4 = DownSample(inplace = 252, outplace = 504, isDown = True) # 32
        
        # self.dilated1 = DilatedConv(inplace = 504, outplace = 128, dilation_rate = 1, stride = 1, kernel_size = 1, padding = 0)
        # self.dilated2 = DilatedConv(inplace = 504, outplace = 128, dilation_rate = 2, stride = 1, kernel_size = 3, padding = 2)
        # self.dilated3 = DilatedConv(inplace = 504, outplace = 128, dilation_rate = 5, stride = 1, kernel_size = 3, padding = 5)
        # self.dilated4 = DilatedConv(inplace = 504, outplace = 128, dilation_rate = 7, stride = 1, kernel_size = 3, padding = 7)
        
        self.layer = nn.Sequential(
            nn.Conv2d(504, 504, (3, 3), stride=1, padding=1, bias=False),
            nn.BatchNorm2d(504),
            nn.ReLU(inplace = True),
        )
        
        self.up_layer1 = UpSample(in_plances = 504, out_plances = 252, isUp = True)  # 64
        self.up_layer2 = UpSample(in_plances = 504, out_plances = 126, isUp = True)   # 128
        self.up_layer3 = UpSample(in_plances = 252, out_plances = 12, isUp = True)   # 256
        self.up_layer4 = UpSample(in_plances = 24, out_plances = 1, isUp = True)     # 512
        
        
        
    
    def forward(self, x):
        layer1 = self.down_layer1(x)
        layer2 = self.down_layer2(layer1)
        layer3 = self.down_layer3(layer2)
        layer4 = self.down_layer4(layer3)
        
        # d1 = self.dilated1(layer4)
        # d2 = self.dilated2(layer4)
        # d3 = self.dilated3(layer4)
        # d4 = self.dilated4(layer4)
        
        # down_last1 = t.cat((d1, d2, d3, d4), 1)
        down_last2 = self.layer(layer4)
        #拼接数据
        # x = t.cat((down_last1, down_last2), 1)
        # 上采样
        x = self.up_layer1(down_last2)
        x = t.cat((x, layer3), 1)
        
        x = self.up_layer2(x)
        x = t.cat((x, layer2), 1)
        
        x = self.up_layer3(x)
        x = t.cat((x, layer1), 1)
        
        x = self.up_layer4(x)
        return x